import json
from typing import Dict, Any

from utils.llm_client import llm_chat
import re
from json_repair import repair_json


def robust_extract_json(text, type='list') -> Any:
    ''' 从文本中提取json数据（列表或字典） '''
    # 如果是列表类型，则左括号为'['，右括号为']'；如果是字典类型，则左括号为'{'，右括号为'}'
    if type == 'list':
        left_bracket, right_bracket = '[', ']'
    else:
        left_bracket, right_bracket = '{', '}'

    stack = []
    start_index = -1
    for i, c in enumerate(text):
        if c == left_bracket:
            if not stack:
                start_index = i  # 记录起始位置
            stack.append(c)
        elif c == right_bracket:
            if stack:
                stack.pop()
                if not stack:  # 找到匹配的最外层括号
                    try:
                        return json.loads(text[start_index:i + 1])
                    except json.JSONDecodeError:
                        continue  # 继续查找其他可能
    return None


def fix_incomplete_json_output(raw_results):
    """ 修复可能不完整的 JSON 输出，确保以 "\n}\n```" 结尾 """
    modified_results = []
    for result in raw_results:
        if isinstance(result, str):  # 修正：添加了双引号
            # 找到最后一个句号（包括中文句号和英文句号）的位置
            match = re.search(r'[.。](?=[^.。]*$)', result)  # 查找最后一个句号
            if match:
                last_period_index = match.start()
                # 将最后一个句号后的内容替换为 "\n}\n```"
                result = result[:last_period_index + 1] + '"\n}\n```'  # 修正：添加了双引号
            else:
                # 如果没有找到句号，直接添加结尾
                result += '"\n}\n```'  # 修正：添加了双引号
        modified_results.append(result)
    return modified_results


def xml_outline_to_recursive(self, outline_xml: str) -> Dict[str, Any]:
    ''' 将xml大纲解析为递归结构 '''
    # 示例文档
    example = {
        "title": "论文标题",
        "content": "",
        "refs": [],
        "children": [
            {
                "title": "1. 第一章",
                "content": "",
                "refs": [],
                "children": [
                    {
                        "title": "1.1 第一节",
                        "content": "",
                        "refs": [],
                        "children": [
                            {
                                "title": "1.1.1 第一小节",
                                "content": "1.1.1小节内容",
                                "refs": ["ref1", "ref2"]
                            }
                        ]
                    },
                    {
                        "title": "1.2 第二节",
                        "content": "1.2节内容",
                        "refs": ["ref4"]
                    }
                ]
            },
            {
                "title": "2. 第二章",
                "content": "第二章内容",
                "refs": [],
                "children": [
                    {
                        "title": "2.1 第一节",
                        "content": "2.1节内容",
                        "refs": ["ref2"]
                    }
                ]
            }
        ]
    }
    # 将xml大纲解析为对象列表格式（包括章节标题及引用文献的id列表）
    trans_list_prompt = "请你将xml大纲直接转化为递归结构（每层都包括“title”、“content”、“refs”和“children”字段，如示例所示。其中全部content的内容都为空；只有最里层的refs内容有值，其它refs都为空列表；title字段中要包括标题序号，且在序号前正确数量的添加“#”）（注意：直接给我生成好的json字典，无需其它任何内容）"
    messages = [
        {"role": "user",
         "content": f"大纲：'''{outline_xml}'''\n，递归结构示例：'''{str(example)}'''，\n{trans_list_prompt}"}
    ]
    response = llm_chat(messages, temperature=0.01, model_name='GLM_4_FLASH')
    res = response.choices[0].message.content
    outline_recursive_dict = self.robust_extract_json(res, type='dict')
    return outline_recursive_dict

if __name__ == '__main__':
    import json

    # 假设这是大模型输出的字符串
    llm_output_str = '''\n{\n"content": "这是第一段。\\n\\n这是第二段。\\n\\n\\n这是第三段。",\n"others": [1, 2, 3]\n}'''

    try:
        # 将 JSON 字符串解析为 Python 字典
        data = json.loads(llm_output_str)

        # 访问 content 字段
        content = data["content"]
        print(content)

        # 你可以根据需要处理 content 中的空行
        # 例如，将多个连续空行替换为一个：
        import re

        normalized_content = re.sub(r'\n\s*\n', '\n\n', content)
        print(normalized_content)


    except json.JSONDecodeError as e:
        print(f"JSON 解析错误: {e}")
        # 在这里处理错误，例如：
        # - 记录错误日志
        # - 尝试修复 JSON 字符串（非常困难，不推荐）
        # - 提示用户或重新生成输出